//
// Created by Tony on 2024/6/11.
//

#include "CodeGenerator.h"
#include "../common/KeyWords.h"
#include "../common/OpCode.h"
#include "../common/Utils.h"
#include "SymbolTable.h"
#include <cstdint>
#include <cstring>
#include <memory>

namespace Compiler {

using Common::OpCode;
using std::make_shared;
using std::static_pointer_cast;

int32_t CodeGenerator::visit(Parser::Ast *ast) { return ast->accept(this); }

int32_t CodeGenerator::visit(Parser::BooleanLiteral *ast) { return 0; }
int32_t CodeGenerator::visit(Parser::CharLiteral *ast) { return 0; }
int32_t CodeGenerator::visit(Parser::Flt64Literal *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::Int64Literal *ast) {
    auto index = constantPool().addInt64(ast->getValue());
    code()->addOpCodeWithUint16(OpCode::LOAD_K, index);
    code()->incTop();
    return 0;
}

int32_t CodeGenerator::visit(Parser::NullLiteral *ast) { return 0; }
int32_t CodeGenerator::visit(Parser::StringLiteral *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::Name *ast) {
    auto symbol = ast->getSymbol();

    if (symbol.scope == ScopeType::CLASS && symbol.type == SymbolType::FIELD) {
        auto field = mClass->getField(ast->name());
        auto fieldRef = constantPool().addFieldRef(mClass->getName(), ast->name());
        if (field->isStatic()) {
            if (ast->isLeftNearbyAssign()) {
                code()->addOpCodeWithUint16(OpCode::PUT_STATIC, fieldRef);
                code()->decTop();
            } else {
                code()->addOpCodeWithUint16(OpCode::GET_STATIC, fieldRef);
                code()->incTop();
            }
            return 0;
        }

        code()->addOpCodeWithUint16(OpCode::LOAD_V, 0); // load this
        code()->incTop();
        if (ast->isLeftNearbyAssign()) {
            code()->addOpCodeWithUint16(OpCode::PUT_FIELD, fieldRef);
            code()->decTop(2);
        } else {
            code()->addOpCodeWithUint16(OpCode::GET_FIELD, fieldRef);
        }
        return 0;
    }

    if (symbol.scope == ScopeType::CLASS && symbol.type == SymbolType::METHOD) {
        auto method = mClass->getMethod(ast->name());
        auto methodRef = constantPool().addMethodRef(mClass->getName(), ast->name());
        if (method->isStatic()) {
            code()->addOpCodeWithUint16(OpCode::INVOKE_STATIC, methodRef);
        } else {
            code()->addOpCodeWithUint16(OpCode::LOAD_V, 0); // load this
            if (method->isPublic()) {
                code()->addOpCodeWithUint16(OpCode::INVOKE_VIRTUAL, methodRef);
            } else {
                code()->addOpCodeWithUint16(OpCode::INVOKE_SPECIAL, methodRef);
            }
        }
        code()->decTop(method->getArgSlotCount());
        if (!method->isReturnVoid()) {
            code()->incTop();
        }
        return 0;
    }

    if (symbol.scope == ScopeType::BLOCK || symbol.scope == ScopeType::METHOD) {
        if (ast->isLeftNearbyAssign()) {
            code()->addOpCodeWithUint16(OpCode::STORE_V, symbol.slotId);
            code()->decTop();
        } else {
            code()->addOpCodeWithUint16(OpCode::LOAD_V, symbol.slotId);
            code()->incTop();
        }
        return 0;
    }

    if (symbol.scope == ScopeType::IMPORTS) {
        // TODO
    }

    return 0;
}

int32_t CodeGenerator::visit(Parser::Operator *ast) {
    const auto *op = ast->op();
    if (strcmp(op, "+") == 0) {
        code()->addOpCode(OpCode::ADD);
    } else if (strcmp(op, "-") == 0) {
        code()->addOpCode(OpCode::SUB);
    } else if (strcmp(op, "*") == 0) {
        code()->addOpCode(OpCode::MUL);
    } else if (strcmp(op, "/") == 0) {
        code()->addOpCode(OpCode::DIV);
    } else if (strcmp(op, "<") == 0) {
        code()->addOpCode(OpCode::LT);
    } else if (strcmp(op, ">") == 0) {
        code()->addOpCode(OpCode::GT);
    }
    code()->decTop();
    return 0;
}

int32_t CodeGenerator::visit(Parser::BinaryExpr *ast) {
    const auto *op = ast->op()->op();

    // assign
    if (strcmp(op, "=") == 0) {
        ast->right()->accept(this);
        ast->left()->accept(this);
        return 0;
    }

    // method invoke
    if (strcmp(op, "(") == 0) {
        ast->right()->accept(this);
        ast->left()->accept(this);
        return 0;
    }

    if (strcmp(op, ".") == 0) { // 当前实现只能处理left是对象的情况，无法处理是类的情况
        auto left = static_pointer_cast<Parser::Name>(ast->left());
        auto leftSymbol = left->getSymbol();
        auto leftClass = mClassTable->getClassById(leftSymbol.classId);

        left->accept(this);

        auto right = static_pointer_cast<Parser::Name>(ast->right());
        auto rightSymbol = right->getSymbol();

        // FIELD
        if (rightSymbol.type == SymbolType::FIELD) {
            auto fieldRef = constantPool().addFieldRef(leftClass->getName(), right->name());
            auto field = leftClass->getField(right->name());
            if (field->isStatic()) {
                if (right->isLeftNearbyAssign()) {
                    code()->addOpCodeWithUint16(OpCode::PUT_STATIC, fieldRef);
                    code()->decTop();
                } else {
                    code()->addOpCodeWithUint16(OpCode::GET_STATIC, fieldRef);
                    code()->incTop();
                }
            } else {
                if (right->isLeftNearbyAssign()) {
                    code()->addOpCodeWithUint16(OpCode::PUT_FIELD, fieldRef);
                    code()->decTop(2);
                } else {
                    code()->addOpCodeWithUint16(OpCode::GET_FIELD, fieldRef);
                }
            }
            return 0;
        }

        // METHOD
        auto methodRef = constantPool().addMethodRef(leftClass->getName(), right->name());
        auto method = leftClass->getMethod(right->name());
        if (method->isStatic()) {
            code()->addOpCodeWithUint16(OpCode::INVOKE_STATIC, methodRef);
        } else if (method->isPublic()) {
            code()->addOpCodeWithUint16(OpCode::INVOKE_VIRTUAL, methodRef);
        } else {
            code()->addOpCodeWithUint16(OpCode::INVOKE_SPECIAL, methodRef);
        }
        code()->decTop(method->getArgSlotCount());
        if (!method->isReturnVoid()) {
            code()->incTop();
        }
        return 0;
    }

    // logic
    if (strcmp(op, "||") == 0) {
        // TODO logic op support shot circuit
        return 0;
    }

    ast->left()->accept(this);
    ast->right()->accept(this);
    ast->op()->accept(this);
    return 0;
}

int32_t CodeGenerator::visit(Parser::TernaryExpr *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::UnaryExpr *ast) {
    const auto *op = ast->op()->op();
    if (strcmp(op, KeyWords::NEW) == 0) {
        auto expr = static_pointer_cast<Parser::BinaryExpr>(ast->expr());
        auto className = static_pointer_cast<Parser::Name>(expr->left())->name();
        auto classRef = constantPool().addClassRef(className);
        code()->addOpCodeWithUint16(OpCode::NEW, classRef);
        auto objSP = code()->top();
        code()->incTop();
        auto args = static_pointer_cast<Parser::ArgumentList>(expr->right());
        args->accept(this);
        auto methodRef = constantPool().addMethodRef(className, Runtime::OBJECT_INIT_METHOD_NAME);
        code()->addOpCodeWithUint16(OpCode::LOAD_V, objSP);
        code()->incTop();
        code()->addOpCodeWithUint16(OpCode::INVOKE_SPECIAL, methodRef);
        code()->decTop(args->numChildren() + 1);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::TypeTag *ast) { return 0; }
int32_t CodeGenerator::visit(Parser::AccessFlags *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::SimpleStmt *ast) {
    ast->getExpr()->accept(this);
    // TODO: add POP to optimize memory usage
    return 0;
}

int32_t CodeGenerator::visit(Parser::AssertStmt *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::BlockStmt *ast) {
    for (auto &child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::ForStmt *ast) {
    ast->getInitExpr()->accept(this);

    auto pcBeginLoop = code()->size();
    ast->getConditionExpr()->accept(this);
    auto pcBreakLoop = code()->addOpCodeWithInt16(OpCode::IF_FALSE, 0);

    ast->getBlock()->accept(this);
    ast->getUpdateExpr()->accept(this);

    code()->addOpCodeWithInt16(OpCode::JMP, backwardOffset(pcBeginLoop));
    code()->updateOperandInt16(pcBreakLoop + 1, forwardOffset(pcBreakLoop));
    return 0;
}

int32_t CodeGenerator::visit(Parser::IfStmt *ast) {
    ast->getCondition()->accept(this);

    auto pcSkipIfBlock = code()->addOpCodeWithInt16(OpCode::IF_FALSE, 0);
    code()->decTop();

    ast->getIfBlock()->accept(this);

    auto elseBlock = ast->getElseBlock();
    if (elseBlock == nullptr) {
        code()->updateOperandInt16(pcSkipIfBlock + 1, forwardOffset(pcSkipIfBlock));
        return 0;
    }
    // if hit if branch and has else branch need skip else block code section
    auto pcSkipElseBlock = code()->addOpCodeWithInt16(OpCode::JMP, 0);

    code()->updateOperandInt16(pcSkipIfBlock + 1, forwardOffset(pcSkipIfBlock));

    elseBlock->accept(this);

    code()->updateOperandInt16(pcSkipElseBlock + 1, forwardOffset(pcSkipElseBlock));
    return 0;
}

int32_t CodeGenerator::visit(Parser::NullStmt *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::VarDefStmt *ast) {
    auto initExpr = ast->getInitExpr();
    CHECK_NULL_RETURN(initExpr, 0);

    initExpr->accept(this);

    code()->addOpCodeWithUint16(OpCode::STORE_V, ast->getSymbol().slotId);
    code()->decTop();
    return 0;
}

int32_t CodeGenerator::visit(Parser::ReturnStmt *ast) {
    auto retExpr = ast->getReturnExpr();
    if (retExpr == nullptr) {
        code()->addOpCode(OpCode::RETURN);
        return 0;
    }

    retExpr->accept(this);
    code()->addOpCode(OpCode::RETURN_V);
    code()->decTop();
    return 0;
}

// array
int32_t CodeGenerator::visit(Parser::ArrayIndex *ast) { return 0; }

// function
int32_t CodeGenerator::visit(Parser::ParameterList *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::FunctionDefine *ast) {
    auto funcSymbolTable = ast->getSymbolTable();

    auto function = make_shared<RtCode::Method>(ast->name()->name());
    function->setLocalSlotCount(funcSymbolTable->maxSlotCount());
    function->addAccessFlag(RtCode::AccessFlag::STATIC);

    auto funcCodeGenerator = make_shared<CodeGenerator>(mClass, mClassTable, function);

    ast->block()->accept(funcCodeGenerator.get());
    return 0;
}

int32_t CodeGenerator::visit(Parser::ArgumentList *ast) {
    for (auto &child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

// module
int32_t CodeGenerator::visit(Parser::Module *ast) {
    for (auto &child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::FieldDefine *ast) {
    auto initExpr = ast->getInitExpr();
    if (initExpr == nullptr) {
        return 0;
    }

    auto fieldName = ast->getName();
    auto field = mClass->getField(fieldName);
    auto fieldRef = constantPool().addFieldRef(mClass->getName(), field->getName());

    // switch current method for code generator
    mMethod = field->isStatic() ? mClass->getClassInitMethod() : mClass->getObjectInitMethod();

    if (field->isStatic()) {
        initExpr->accept(this);
        code()->addOpCodeWithUint16(OpCode::PUT_STATIC, fieldRef);
        code()->decTop();
    } else {
        initExpr->accept(this);
        code()->addOpCodeWithUint16(OpCode::LOAD_V, 0);
        code()->incTop();
        code()->addOpCodeWithUint16(OpCode::PUT_FIELD, fieldRef);
        code()->decTop(2);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::MethodDefine *ast) {
    auto methodName = ast->getName();
    mMethod = mClass->getMethod(methodName);
    ast->getBody()->accept(this);
    if (mMethod->isReturnVoid() && !code()->hasReturn()) {
        code()->addOpCode(OpCode::RETURN);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::NativeMethodDecl *ast) {
    mMethod = mClass->getMethod(ast->getName());
    if (!mMethod->isReturnVoid()) {
        code()->incTop();
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::ClassBody *ast) {
    for (auto &child : ast->getChildren()) {
        child->accept(this);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::ClassDefine *ast) {
    mClass = mClassTable->getClassById(ast->getClassId());
    ast->getClassBody()->accept(this);

    mClass->getClassInitMethod()->getCode()->addOpCode(OpCode::RETURN);

    mClass->getObjectInitMethod()->getCode()->addOpCode(OpCode::RETURN);
    mClass->getObjectInitMethod()->setLocalSlotCount(1);
    return 0;
}

int32_t CodeGenerator::visit(Parser::ClassModule *ast) {
    for (auto &classDefine : ast->getClassDefines()) {
        classDefine->accept(this);
    }
    return 0;
}

int32_t CodeGenerator::visit(Parser::ViewBuilder *ast) {
    // calc args
    ast->getArgumentList()->accept(this);

    // new object
    auto className = ast->getName()->name();
    auto ownerClass = mClassTable->getClassByName(className);
    auto methodRef = constantPool().addMethodRef(className, "create");
    code()->addOpCodeWithUint16(OpCode::INVOKE_STATIC, methodRef);
    auto method = ownerClass->getMethod("create");
    code()->decTop(method->getArgSlotCount()); // method args

    auto objSP = code()->top();
    code()->incTop(); // method return

    // set attribute
    for (int i = 0; i < ast->getAttributeSetterList()->numChildren(); i++) {
        auto attributeSetter = ast->getAttributeSetterList()->getAttributeSetter(i);
        auto args = attributeSetter->getArgumentList();
        args->accept(this);
        code()->addOpCodeWithUint16(OpCode::LOAD_V, objSP);
        code()->incTop();
        auto methodRef = constantPool().addMethodRef(className, attributeSetter->getName()->name());
        code()->addOpCodeWithUint16(OpCode::INVOKE_VIRTUAL, methodRef);
        code()->decTop(args->numChildren() + 1);
    }

    // add children
    auto viewBuilderList = ast->getViewBuilderList();
    CHECK_NULL_RETURN(viewBuilderList, 0);
    for (int i = 0; i < viewBuilderList->numChildren(); i++) {
        auto childViewBuilder = viewBuilderList->getViewBuilder(i);
        childViewBuilder->accept(this);
        code()->addOpCodeWithUint16(OpCode::LOAD_V, objSP);
        code()->incTop();
        auto methodRef = constantPool().addMethodRef(className, "addChild");
        code()->addOpCodeWithUint16(OpCode::INVOKE_VIRTUAL, methodRef);
        code()->decTop(2);
    }

    return 0;
}

int32_t CodeGenerator::visit(Parser::ViewBuilderList *ast) { return 0; }

int32_t CodeGenerator::visit(Parser::ViewBuilderBlock *ast) {
    ast->getViewBuilder()->accept(this);
    code()->addOpCode(OpCode::RETURN_V);
    return 0;
}

int32_t CodeGenerator::visit(Parser::AttributeSetter *ast) { return 0; }
int32_t CodeGenerator::visit(Parser::AttributeSetterList *ast) { return 0; }

} // namespace Compiler
